package com.timeriver.feature_project

import org.apache.spark.ml.feature.Bucketizer
import org.apache.spark.sql.{DataFrame, SparkSession}

/**
  * Bucketizer将一列连续特征转换为一列特征桶，其中的桶由用户指定。它需要的参数如下：
  *   split：用于将连续特征映射到桶中的参数。有了n+1个分片，就有了n个桶。一个由splits x,y定义的桶可以容纳[x,y]范围内的值，
  *         除了最后一个桶，它也包括y。在-inf，inf的值必须明确提供，以涵盖所有的Double值；否则，在指定的分片之外的值将被当
  *         作错误处理。
  *   涵盖全范围的例子：Array(Double.NegativeInfinity, 0.0, 1.0, Double.PositiveInfinity)
  */
object BucketizerDemo {
  def main(args: Array[String]): Unit = {
    val session: SparkSession = SparkSession.builder()
      .master("local[6]")
      .appName("连续特征给定边界离散化")
      .getOrCreate()

    val ints: Array[Int] = Array(95, 45, 56, 89, 21, 5, 12, 62, 32)

    val data: DataFrame = session.createDataFrame(ints.map(Tuple1.apply)).toDF("label")

    /** 将0-100以内的数字分为5个桶 */
    val splits = Array(0.0, 20.0, 40.0, 60.0, 80.0, 100.0)

    val bucketizer = new Bucketizer()
      .setInputCol("label")
      .setOutputCol("bucketedLabel")
      .setSplits(splits)

    val res: DataFrame = bucketizer.transform(data)
    res.show(false)
  }
}
